numpy array를 통해 매트릭스 및 텐서 연산을 할때, 연산의 규칙이 꽤나 복잡하다.
넘파이에서 제공하는 내적연산에는 dot과 matmul이 있다.
dot은 벡터를 대상으로 하는 내적 연산이기 때문에 다차원 매트릭스에 대해 dot 연산을 수행하면, 의도했던 내적연산이 발생하지 않는다.
매트릭스와 텐서를 대상으로 하는 내적 연산은 matmul이다.

텐서에서 두 연산자의 주요한 차이를 이해하기 위해, 먼저 브로드캐스팅의 개념을 이해해야한다.

브로드캐스팅Broadcasting

numpy array에서 벡터 및 매트릭스 연산에서 제공되는 브로드캐스팅이 있다.
브로드캐스팅의 목적은 차원이 다른 두 텐서(벡터나 매트릭스도 포함)간의 연산을 지원하는 데에 있다.
구체적으론 두 스텝을 통해 브로드캐스팅이 이뤄진다.

1. 차원확장Expand dimensions
2. 크기정렬Align Shapes

Rule 1: 차원확장Expand dimensions

차원확장은 차원의 개수가 다를때, 1을 추가하여 개수를 맞춰주는 것이다.
eg.) shape (2, 3, 4, 5, 6)과 shape (7, 8, 9)이 있을 때, 각각 차원의 개수가 5개, 3개로 다르다. 따라서 shape (7, 8, 9)를 shape (1, 1, 7, 8, 9)로 확장하여 차원수를 맞춰준다.
*numpy array로 보면 다음과 같이 변하는 것이다 [...] -> [[[...]]]

Rule 2: 크기정렬Align Shapes

크기정렬은 차원의 수가 같다고 연산이 가능해지는 것이 아니기 때문에 존재한다.
차원의 사이즈도 맞춰줘야한다. 즉, shape가 같도록 만들어줘야 한다.
eg.) shape (2, 3, 4, 5, 6)과 shape (1, 1, 7, 8, 9)의 shape가 같아지려면 사이즈가 1인 차원들을 각각 2, 3으로 복제를 통해 정렬할 수 있다. 문제는 여기서 뒤 3개의 차원은 정렬이 불가능하다는 것이다. 복제를 통해 맞추는 것이 불가능하기 때문이다. (차원이 이미 여러개인데 무엇을 복제하겠는가?) 따라서 이 예시는 최종적으로 브로드캐스팅에 실패하고 shape mismatch 에러를 발생시킬 것이다.

따라서 align이 가능하려면 각 차원의 사이즈가 1 또는 서로 같아야 할것이다.

ex

A = np.random.rand(2, 3, 4, 5, 6) # shape: (2, 3, 4, 5, 6)
B = np.random.rand(3, 4, 6, 7) # shape: (3, 4, 6, 7)

# ex1
(A[:3] + B[:2]).shape # (2, 3, 4) + (3, 4) -> (2, 3, 4) + (2, 3, 4)
# 덧셈연산에서는 뒤의 두개의 차원을 포함하면 브로드캐스팅 에러가 발생한다.
>>> (2, 3, 4)

# ex2
(A @ B).shape # @은 matmul연산자
>>> (2, 3, 4, 5, 7)

ex1)
(2, 3, 4)와 (3, 4)에 대해 덧셈 연산이 불가능하기 때문에 먼저 브로드캐스팅을 수행하여 B: (3, 4) -> (1, 3, 4) -> (2, 3, 4)로 서로 동일한 차원 사이즈를 맞주고 덧셈을 수행한다.

ex2)
matmul연산은 마지막 두 차원에서 매트릭스 내적을 하고 나머지 차원에서만 브로드캐스트 sum을 수행한다. 따라서 (5, 6)과 (6, 7)은 내적을 수행하여 (5, 7)의 shape가 되고 앞 차원은 ex1)과 동일하게 브로드캐스팅하여 합을 해준다. 이 부분은 후술할 matmul 연산의 메커니즘을 먼저 이해하면 된다.

dot vs matmul

numpy에서 dot 연산과 matmul 연산에는 한가지 중요한 차이가 있다.
바로 브로드캐스팅의 여부이다.

하지만 깊은 이해를 하기 위해서는 수식적 정의로 접근해야한다.
그 전에 우리가 사용하는 내적이라는 용어에 대해 수식적인 정의에서 시작하겠다.

내적

일반적으로 사용되는 정의는 벡터와 벡터를 곱하여 스칼라 결과를 얻는 것이다.
하지만 여기서 우리는 용어를 조금 확장해서 element-wise prodcut를 가리키는 용어로 사용한다.

결국 벡터든 매트릭스든 텐서든간에 차원이라는 것은 원소들로 이루어져 있다. 예컨대 1부터 100까지의 수 100개를 각각 원소로 보고 이를 A = reshape(10, 10, 10)하면, 3개의 차원으로 정렬한 셈이 되는 것이다. 그렇다면 각각의 원소는 $A_{ijk}$로 표현된다. 이때 B = A.copy()로 정의하자. 이때 A와 B의 내적은 다음 중 무엇일까?

셋 모두이다. 물론 (1)의 결과는 스칼라이고, (2)의 결과는 벡터이고, (3)의 결과는 매트릭스로 결과의 차원은 모두 다르다. 하지만 이들은 모두 내적이다. 즉, 우리가 어떻게 내적하느냐의 따라 연산의 결과가 천차만별로 달라질 수 있다는 것이다.

거기서 dot과 matmul의 차이점이 발생한다.

텐서 내적의 수식적 예시

(1) dot
A의 shape: (i, j, k, p)
B의 shape: (m, p, n)
dot(A, B)의 shape: (i, j, k, m, n)

(2) matmul
A의 shape: (x, y, z, k, p)
B의 shape: (x, y, z, p, n)
matmul(A, B)의 shape: (x, y, z, k, n)

이들은 모두 공통적으로 A[-1]과 B[-2]의 차원에 대해 곱을 수행하며, 나머지 차원에 대해 다루는 방식에서 차이가 발생한다.
여기서 matmul연산의 경우 마지막 두 차원은 서로 매트릭스 내적으로 처리하고 그 외의 차원에 대해서만 일치하면 된다(dot은 매트릭스 내적으로 생각하면 차원의 순서가 꼬임에 주의).
내가 배울때는 두 연산자의 차원 관점에 대해서도 배웠다. 하지만 수식적 서술에 의하면 그러한 관점이 완전히 불필요한 설명이라는 생각이 들어 생략한다.

dot과 matmul은 이것으로 끝이다. 다만, 여기서 matmul은 브로드캐스팅이 적용되는 연산이다. 앞부분 차원의 shape이 서로 다르더라도 브로드캐스팅만 가능하다면 연산이 이뤄진다.

matmul연산의 일반화는 코드 구현 참조.

matmul의 코드 구현

def generalized_matmul(A, B):
    """
    브로드캐스팅을 포함한 행렬 곱셈을 수행하는 함수.
    
    Parameters:
    A : ndarray
        좌측 텐서. 형식은 (..., k, l)
    B : ndarray
        우측 텐서. 형식은 (..., l, n)
        
    Returns:
    result : ndarray
        브로드캐스팅이 포함된 행렬 곱셈의 결과. 형식은 (..., k, n)
    """
    # 마지막 두 차원이 행렬 곱셈이 가능한지 검사
    if A.shape[-1] != B.shape[-2]:
        raise ValueError("행렬 곱셈을 위해 A의 마지막 차원과 B의 마지막에서 두 번째 차원이 일치해야 합니다.")
    
    # 브로드캐스트될 형식 결정
    broadcast_shape = np.broadcast_shapes(A.shape[:-2], B.shape[:-2])
    
    # A와 B를 브로드캐스트 형식으로 맞춰 재배열
    A_broadcasted = np.broadcast_to(A, broadcast_shape + A.shape[-2:])
    B_broadcasted = np.broadcast_to(B, broadcast_shape + B.shape[-2:])
    
    # 마지막 두 축을 기준으로 행렬 곱셈 수행
    result = np.einsum('...kl,...ln->...kn', A_broadcasted, B_broadcasted)
    
    return result

# 예제 사용법
A = np.random.rand(2, 3, 4, 5, 6)  # 형식: (2, 3, 4, 5, 6)
B = np.random.rand(1, 3, 4, 6, 7)  # 형식: (1, 3, 4, 6, 7)

# 앞쪽 차원에 대해 브로드캐스팅을 포함한 결과 계산
C = generalized_matmul(A, B)
print("결과 형식:", C.shape)  # 예상 형식: (2, 3, 4, 5, 7)